package com.songlanyun.common.streamquery;

import cn.hutool.core.collection.CollectionUtil;
import cn.hutool.core.util.ObjectUtil;
import org.apache.ibatis.session.ResultContext;
import org.apache.ibatis.session.ResultHandler;

import java.util.ArrayList;
import java.util.List;
import java.util.function.Consumer;

public  class QueryStreamHandLer<T> implements ResultHandler<T> {

    protected int batchSize = 100;
    private int size = 0;
    private int batchNum = 0;

    public Consumer<BaseBatchData> consumer;

    public QueryStreamHandLer() {

    }

    public QueryStreamHandLer(int batchSize) {
        this.batchSize = batchSize;
    }

    /**
     * 存储每批数据的临时容器
     */
    private List<T> batchResult = new ArrayList<>();

    @Override
    public void handleResult(ResultContext<? extends T> resultContext) {
        T row = resultContext.getResultObject();
        // 你可以看自己的项目需要分批进行处理或者单个处理，这里以分批处理为例
        batchResult.add(row);
        size++;
        if (size == this.batchSize) {
            handleProcess();
        }
    }


    private void handleProcess() {
        batchNum++;
        try {
            // 在这里可以对你获取到的批量结果数据进行需要的业务处理
            if (CollectionUtil.isNotEmpty(batchResult)) {
                if (ObjectUtil.isNotNull(consumer)) {
                    consumer.accept(new BaseBatchData() {
                        {
                            setBatchNum(batchNum);
                            setData(batchResult);
                        }
                    });
                }
            }
        } finally {
            // 处理完每批数据后后将临时清空
            size = 0;
            batchResult.clear();
        }
    }

    public void endProcess() {
        handleProcess();// 处理最后一批不到BATCH_SIZE的数据
    }
}
